# coding=utf-8
# 用opencv实现兑换站识别
import cv2
import numpy as np
import pyrealsense2 as rs
import time
#ros发布的部分
import rospy
from geometry_msgs.msg import Point, Vector3

# 求取图像上一个点的三维坐标
def get_3d_point(depth_frame, xy):
    x, y = int(xy[0]), int(xy[1])
    # 获取像素点的深度值
    depth = depth_frame.get_distance(x, y)
    # 如果是噪点，则用四周不为0的点的深度值的平均值代替
    if depth < 0.01:
        count = 0
        for i in range(-3, 4):
            for j in range(-3, 4):
                # 防止越界
                x_limit = max(min(x+i, 639), 0)
                y_limit = max(min(y+j, 479), 0)
                if depth_frame.get_distance(x_limit, y_limit) > 0.01:
                    depth += depth_frame.get_distance(x_limit, y_limit)
                    count += 1
        if count != 0:
            depth = depth / count
    if depth == 0:
        print("depth is 0")
    # intelrealsense参数
    intrinsics = depth_frame.profile.as_video_stream_profile().intrinsics
    # 像素坐标系转换为相机坐标系
    camera_coordinate = rs.rs2_deproject_pixel_to_point(intrinsics, [x, y], depth)
    return camera_coordinate

# 三点坐标 -> 正方体中心坐标，正方体三个互相垂直面的单位法向量(第一个法向量指向z轴正向)
def getCenter(p1, p2, p3, roll_45):
    p1 = np.array(p1)
    p2 = np.array(p2)
    p3 = np.array(p3)
    # 化p1,p3在同一条对角线上
    p1p2 = np.linalg.norm(p1 - p2)
    p1p3 = np.linalg.norm(p1 - p3)
    p2p3 = np.linalg.norm(p2 - p3)
    if(p1p2 > p1p3 and p1p2 > p2p3):
        p3, p2 = p2, p3
    elif(p2p3 > p1p2 and p2p3 > p1p3):
        p1, p2 = p2, p1
    # p1p2向量(一个法向量)
    vec2 = np.array(p2) - np.array(p1)
    # p2p3向量(另一个法向量)
    vec3 = np.array(p3) - np.array(p2)
    # 正面法向量
    vec1 = np.cross(vec2, vec3)
    # print("vec1: ", vec1)
    # 法向量单位化
    vec_norm1 = np.linalg.norm(vec1)
    unit_vec1 = vec1 / (vec_norm1 if vec_norm1 != 0 else 1)
    vec_norm2 = np.linalg.norm(vec2)
    unit_vec2 = vec2 / (vec_norm2 if vec_norm2 != 0 else 1)
    vec_norm3 = np.linalg.norm(vec3)
    unit_vec3 = vec3 / (vec_norm3 if vec_norm3 != 0 else 1)
    # 旋转45度
    if roll_45:
        unit_vec2 = (unit_vec2 + unit_vec3) / 2 ** 0.5
        unit_vec3 = (unit_vec2 - unit_vec3) / 2 ** 0.5
    # z轴法向量指向正方向
    if(unit_vec1[0] < 0):
        unit_vec1 = -unit_vec1
    # 平面中心
    front_center = (p1 + p3) / 2
    # 固定棱长(m)
    gold_a = 0.24
    # 计算棱长
    # gold_a = np.linalg.norm(p1 - p2)
    # 中心坐标
    center = front_center + gold_a / 2 * unit_vec1
    # t=体心坐标，z轴方向单位法向量，另外两个方向单位向量
    return center, unit_vec1, unit_vec2, unit_vec3

# 输出结果
def printResult(depth_frame, points, roll_45):
    # 绘制points中的点
    for j in range(len(points)):
        cv2.circle(color_image, (points[j][0], points[j][1]), 4, (255, 255, 255), -1)
    # 读取points中的点的三维坐标
    points_3d = []
    for i in range(3):
        dep = get_3d_point(depth_frame, points[i])
        points_3d.append(dep)
    # 输出结果
    if len(points_3d) >= 3:
        output = getCenter(points_3d[0], points_3d[1], points_3d[2], roll_45)
        print("center: ", output[0])
        print(output[1:4])
    print("")

# 两点距离
def getDistance(p1, p2):
    p1 = np.array(p1)
    p2 = np.array(p2)
    return np.linalg.norm(p1 - p2)

#ROS_publish_topic,main中需要调用此函数
def publish_geometry(center, unit_vec1, unit_vec2, unit_vec3):
    rospy.init_node('publish_geometry', anonymous=True)
    rate = rospy.Rate(10) # 10hz

    # 创建四个消息发布者
    center_pub = rospy.Publisher('center', Point, queue_size=10)
    unit_vec1_pub = rospy.Publisher('unit_vec1', Vector3, queue_size=10)
    unit_vec2_pub = rospy.Publisher('unit_vec2', Vector3, queue_size=10)
    unit_vec3_pub = rospy.Publisher('unit_vec3', Vector3, queue_size=10)

    while not rospy.is_shutdown():
        # 发布center
        center_msg = Point()
        center_msg.x = center[0]
        center_msg.y = center[1]
        center_msg.z = center[2]
        center_pub.publish(center_msg)

        # 发布unit_vec1
        unit_vec1_msg = Vector3()
        unit_vec1_msg.x = unit_vec1[0]
        unit_vec1_msg.y = unit_vec1[1]
        unit_vec1_msg.z = unit_vec1[2]
        unit_vec1_pub.publish(unit_vec1_msg)

        # 发布unit_vec2
        unit_vec2_msg = Vector3()
        unit_vec2_msg.x = unit_vec2[0]
        unit_vec2_msg.y = unit_vec2[1]
        unit_vec2_msg.z = unit_vec2[2]
        unit_vec2_pub.publish(unit_vec2_msg)

        # 发布unit_vec3
        unit_vec3_msg = Vector3()
        unit_vec3_msg.x = unit_vec3[0]
        unit_vec3_msg.y = unit_vec3[1]
        unit_vec3_msg.z = unit_vec3[2]
        unit_vec3_pub.publish(unit_vec3_msg)

        rate.sleep()

if __name__ == '__main__':

    # 参数设置
    W, H = 640, 480 # 图像宽高
    FPS = 30 # 帧率
    MAX_LENGTH = 1000 # 区分正面和侧面轮廓所使用的空间长度
    RED = (0, 0, 255) # 红色
    BLUE = (255, 0, 0) # 蓝色
    GREEN = (0, 255, 0) # 绿色

    # 设置hsv阈值
    lower_blue = np.array([100, 43, 46])
    upper_blue = np.array([124, 255, 255])
    lower_red = np.array([0, 33, 10])
    upper_red = np.array([10, 255, 255])
    # 创建一个管道
    pipeline = rs.pipeline() 
    # 配置要流式传输的管道
    config = rs.config()
    config.enable_stream(rs.stream.color, W, H, rs.format.bgr8, 30) # 8位rgb图像
    config.enable_stream(rs.stream.depth, W, H, rs.format.z16, 30) # 16位深度图像
    # 左右双目
    config.enable_stream(rs.stream.infrared, 1, W, H, rs.format.y8, 30) # 8位左红外图像
    config.enable_stream(rs.stream.infrared, 2, W, H, rs.format.y8, 30) # 8位右红外图像
    align = rs.align(rs.stream.color)
    # 视野大，但有残缺
    # align = rs.align(rs.stream.depth)
    # 开始传输帧
    profile = pipeline.start(config)
    # 更改rgb相机的曝光值
    color_sensor = profile.get_device().first_color_sensor()
    # 关闭自动曝光
    color_sensor.set_option(rs.option.enable_auto_exposure, 0)
    # 设置曝光值
    color_sensor.set_option(rs.option.exposure, 25)
    # 
    depth_sensor = profile.get_device().first_depth_sensor()
    # 
    depth_scale = depth_sensor.get_depth_scale()

    # 对视频进行循环处理
    while True:
        # 等待开启通道
        frames = pipeline.wait_for_frames()
        # 将深度图和RGB图对齐
        frames = align.process(frames)
        # 获得RGB图像
        color_frame = frames.get_color_frame()
        # 未获得RGB图像
        if not color_frame:
            continue
        # 获得深度图像
        depth_frame = frames.get_depth_frame()
        # 把图像像素转化为数组
        color_image = np.asanyarray(color_frame.get_data())
        depth_image = np.asanyarray(depth_frame.get_data())
        # 深度图像转化为伪彩色图像(alpha数值越大对比度色差越大)
        depth_colormap = cv2.applyColorMap(cv2.convertScaleAbs(depth_image, alpha=0.21), cv2.COLORMAP_JET)
        # 转换到HSV (颜色空间转换) (色调H，饱和度S，亮度V)
        hsv = cv2.cvtColor(color_image, cv2.COLOR_BGR2HSV)
        # 根据阈值构建掩模 (将在两个阈值内的像素值设置为白色，而不在阈值区间内的像素值设置为黑色)
        binary = cv2.inRange(hsv, lower_red, upper_red)
        # 形态学处理
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
        binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel)
        # binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
        # 放慢视频播放速度
        # time.sleep(0.2)
        # 二值化图像
        cv2.imshow('binary', binary)

        # 寻找轮廓
        contours, hierarchy = cv2.findContours(binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        # 正面轮廓
        front_contours = []
        # 侧面轮廓
        side_contours = []
        # num = 0

        # 遍历轮廓分类
        for i in range(len(contours)):
            # 计算轮廓面积
            area = cv2.contourArea(contours[i])
            # print (area)
            # 轮廓外接矩形
            rect = cv2.minAreaRect(contours[i])            
            # 轮廓中最远的两个点
            max_l = 0
            max_idx = [0, 1]
            for j in range(len(contours[i])):
                for k in range(j + 1, len(contours[i])):
                    l = (contours[i][j][0][0] - contours[i][k][0][0]) ** 2 + \
                        (contours[i][j][0][1] - contours[i][k][0][1]) ** 2
                    if l > max_l:
                        max_l = l
                        max_idx = [j, k]
            # 根据轮廓中最远的两个点的空间距离区分轮廓类别
            if max_l > MAX_LENGTH:
                # print("max_l: ",max_l)
                box = cv2.boxPoints(rect)
                box = np.int0(box)
                point_3d1 = get_3d_point(depth_frame, contours[i][max_idx[0]][0])
                point_3d2 = get_3d_point(depth_frame, contours[i][max_idx[1]][0])
                # 两点空间距离
                distance = getDistance(point_3d1, point_3d2)
                # print("distance: ", distance)
                # 正面轮廓
                if(distance < 0.10 or distance > 0.30):
                    cv2.drawContours(color_image, [box], 0, BLUE, 2)
                    front_contours.append(contours[i])
                # 侧面轮廓
                else:
                    cv2.drawContours(color_image, [box], 0, GREEN, 2)
                    side_contours.append(contours[i])

        ## 正面轮廓处理
        if len(front_contours) == 4 or len(front_contours) == 3:
            # 取四个矩形的中心点为正面的中心点
            center = []
            if len(front_contours) == 4:
                for i in range(4):
                    rect = cv2.minAreaRect(front_contours[i])
                    center.append(rect[0])
                # 计算中心点的坐标
                # x, y = (center[0] + center[1] + center[2] + center[3]) / 4 # ???
                x = (center[0][0] + center[1][0] + center[2][0] + center[3][0]) / 4
                y = (center[0][1] + center[1][1] + center[2][1] + center[3][1]) / 4
            elif len(front_contours) == 3:
                center_3d = []
                for i in range(3):
                    rect = cv2.minAreaRect(front_contours[i])
                    center.append(rect[0])
                    center_3d.append(get_3d_point(depth_frame, rect[0]))
                # 空间最长边两端点
                max_length = 0
                max_index = [0, 1]
                for i in range(3):
                    for j in range(i + 1, 3):
                        if getDistance(center_3d[i], center_3d[j]) > max_length:
                            max_length = getDistance(center_3d[i], center_3d[j])
                            max_index = [i, j]
                # 计算中心点的坐标
                # x, y = (center[max_index[0]] + center[max_index[1]]) / 2 # ???
                x = (center[max_index[0]][0] + center[max_index[1]][0]) / 2
                y = (center[max_index[0]][1] + center[max_index[1]][1]) / 2
            # 绘制中心点
            cv2.circle(color_image, (int(x), int(y)), 5, RED, -1)

            ## good_contours保留三个外接矩形面积最大的轮廓
            while len(front_contours) > 3:
                min_area = 100000
                min_index = 0
                for j in range(len(front_contours)):
                    rect = cv2.minAreaRect(front_contours[j])
                    area = rect[1][0] * rect[1][1]
                    if area < min_area:
                        min_area = area
                        min_index = j
                front_contours.pop(min_index)
            
            ## 对每个good_contours中的点进行遍历,找到距离中心点最远的点，放入points中
            points = []
            for i in range(3):
                max_l = 0
                max_point = [0, 0]
                for j in range(len(front_contours[i])):
                    l = getDistance(front_contours[i][j][0], (x, y))
                    if l > max_l:
                        max_l = l
                        max_point = front_contours[i][j][0]
                points.append(max_point)
            # 将points中的点向中心点移动6%并整数化
            for i in range(3):
                points[i][0] = int(points[i][0] + (x - points[i][0]) * 0.06)
                points[i][1] = int(points[i][1] + (y - points[i][1]) * 0.06)
            printResult(depth_frame, points, True)

        ## 侧面轮廓处理
        if len(side_contours) > 0:
            for i in range(len(side_contours)):
                big_rect = cv2.minAreaRect(side_contours[i])
                box = cv2.boxPoints(big_rect)
                box = np.int0(box)
                center = (int(big_rect[0][0]), int(big_rect[0][1]))
                points = []
                # 最下点
                max_y = 0
                max_index = 0
                for j in range(len(side_contours[i])):
                    if side_contours[i][j][0][1] > max_y:
                        max_y = side_contours[i][j][0][1]
                        max_index = j
                points.append(side_contours[i][max_index][0])
                # 距下点最远的点
                max_l = 0
                max_index = 0
                for j in range(len(side_contours[i])):
                    l = getDistance(side_contours[i][j][0], points[0])
                    if l > max_l:
                        max_l = l
                        max_index = j
                points.append(side_contours[i][max_index][0])
                # 距两点距离之和最远的点
                max_l = 0
                max_index = 0
                for j in range(len(side_contours[i])):
                    l = getDistance(side_contours[i][j][0], points[0]) + getDistance(side_contours[i][j][0], points[1])
                    if l > max_l:
                        max_l = l
                        max_index = j
                points.append(side_contours[i][max_index][0])
                # 输出结果
                printResult(depth_frame, points, True)

        ## 识别失败
        if len(front_contours) > 4:
            print("too many contours")
        elif len(front_contours) < 3 and len(side_contours) == 0:
            print("too few contours")

        # 显示图像
        cv2.imshow('color_image', color_image)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cv2.destroyAllWindows()
